package spring.cloud.tasks.executor.name_space;

import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.cache.NodeCache;
import org.apache.curator.framework.recipes.cache.NodeCacheListener;
import org.apache.curator.framework.recipes.cache.TreeCache;
import org.apache.curator.framework.recipes.cache.TreeCacheListener;
import org.apache.curator.utils.CloseableExecutorService;
import spring.cloud.tasks.common.ThreadFactory;
import spring.cloud.tasks.executor.utils.LogUtils;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Slf4j
@Setter
@Getter
public class TaskExecutorNodeCacheOrTreeCacheManager {
    private Map<String/* path */, NodeCache> pathToNodeCacheMap = new HashMap<>();
    private Map<String/* path-depth */, TreeCache> pathAndDepthToTreeCacheMap = new HashMap<>();
    private String executorId;
    private String taskId;
    private CuratorFramework curatorFramework;
    private ExecutorService executorService;

    public TaskExecutorNodeCacheOrTreeCacheManager(CuratorFramework curatorFramework, String executorId, String taskId) {
        this.executorId = executorId;
        this.taskId = taskId;
        this.curatorFramework = curatorFramework;
        this.executorService = Executors.newSingleThreadExecutor(new ThreadFactory(executorId + "-" + taskId + "-watcher", false));
        LogUtils.info(log, taskId, "ZkCacheManager for executor:{} - job:{} created.", executorId, taskId);
    }

    private static String buildPathAndDepthKey(String path, int depth) {
        return path + "-" + depth;
    }

    //
    public void shutdown() {
        // close TreeCache, CloseableExecutorService will be closed, although the pool will not be shutdownNow, but their listeners will accept the interrupt signal
        // so be careful, such as JobDeleteListener, delete job thread cannot be interrupted
        closeTreeCaches();//closeTreeCaches ...closeNodeCaches
        closeNodeCaches();
        if (executorService != null && !executorService.isShutdown()) {
            executorService.shutdown();
        }
    }

    //
    public void addNodeCacheListener(final String path, final NodeCacheListener nodeCacheListener) {
        NodeCache nodeCache = buildNodeCache(path);
        if (nodeCache != null) {
            nodeCache.getListenable().addListener(nodeCacheListener);
        }
    }

    public NodeCache buildNodeCache(String path) {
        try {
            NodeCache nodeCache = pathToNodeCacheMap.get(path);
            if (nodeCache == null) {
                nodeCache = new NodeCache(curatorFramework, path);
                pathToNodeCacheMap.put(path, nodeCache);
                nodeCache.start();
                LogUtils.info(log, taskId, "{} - {} builds nodeCache for path = {}", executorId, taskId, path);
            }
            return nodeCache;
        } catch (Exception e) {
            LogUtils.error(log, taskId, "{} - {}  fails in building nodeCache for path = {}, saturn will not work correctly.", executorId, taskId, path);
            LogUtils.error(log, taskId, e.getMessage(), e);
        }
        return null;
    }

    public void closeNodeCache(String path) {
        NodeCache nodeCache = pathToNodeCacheMap.get(path);
        if (nodeCache != null) {
            try {
                nodeCache.close();
                pathToNodeCacheMap.remove(path);
                LogUtils.info(log, taskId, "{} - {} closed nodeCache, path is {}", executorId, taskId, path);
            } catch (Exception e) {
                LogUtils.error(log, taskId, e.getMessage(), e);
            }
        }
    }

    public void closeNodeCaches() {
        Iterator<Entry<String, NodeCache>> iterator = pathToNodeCacheMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Entry<String, NodeCache> next = iterator.next();
            String path = next.getKey();
            NodeCache nodeCache = next.getValue();
            try {
                nodeCache.close();
                iterator.remove();
                LogUtils.info(log, taskId, "{} - {} closed nodeCache, path is {}", executorId, taskId, path);
            } catch (Exception e) {
                LogUtils.error(log, taskId, e.getMessage(), e);
            }
        }
    }

    //
    public void addTreeCacheListener(final String path, final int depth, final TreeCacheListener treeCacheListener) {
        TreeCache treeCache = buildTreeCache(path, depth);
        if (treeCache != null) {
            treeCache.getListenable().addListener(treeCacheListener);
        }
    }

    private TreeCache buildTreeCache(String path, int depth) {
        try {
            String pathAndDepth = buildPathAndDepthKey(path, depth);
            TreeCache treeCache = pathAndDepthToTreeCacheMap.get(pathAndDepth);
            if (treeCache == null) {
                treeCache = TreeCache.newBuilder(curatorFramework, path).setMaxDepth(depth).setExecutor(new CloseableExecutorService(executorService, false)).build();
                pathAndDepthToTreeCacheMap.put(pathAndDepth, treeCache);
                treeCache.start();
                LogUtils.info(log, taskId, "{} - {}  builds treeCache for path = {}, depth = {}", executorId, taskId, path, depth);
            }
            return treeCache;
        } catch (Exception e) {
            LogUtils.error(log, taskId, "{} - {} fails in building treeCache for path = {}, depth = {}, saturn will not work correctly.", executorId, taskId, path, depth);
            LogUtils.error(log, taskId, e.getMessage(), e);
        }
        return null;
    }

    public void closeTreeCache(String path, int depth) {
        String pathAndDepth = buildPathAndDepthKey(path, depth);
        TreeCache treeCache = pathAndDepthToTreeCacheMap.get(pathAndDepth);
        if (treeCache != null) {
            try {
                treeCache.close();
                pathAndDepthToTreeCacheMap.remove(pathAndDepth);
                LogUtils.info(log, taskId, "{} - {} closed treeCache, path and depth is {}", executorId, taskId, pathAndDepth);
            } catch (Exception e) {
                LogUtils.error(log, taskId, e.getMessage(), e);
            }
        }
    }

    public void closeTreeCaches() {
        Iterator<Entry<String, TreeCache>> iterator = pathAndDepthToTreeCacheMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Entry<String, TreeCache> next = iterator.next();
            String path = next.getKey();
            TreeCache treeCache = next.getValue();
            try {
                treeCache.close();
                iterator.remove();
                LogUtils.info(log, taskId, "{} - {} closed treeCache, path and depth is {}", executorId, taskId, path);
            } catch (Exception e) {
                LogUtils.error(log, taskId, e.getMessage(), e);
            }
        }
    }
    //
}
